/*
  Author: Chen Haodong (shinkunyachen@outlook.com)
  -----
  Copyright (c)2021 - present Chen Haodong. All Rights Reserved.
 */

#include "checkers/BranchChecker/BranchChecker.h"
#include "framework/Common.h"
extern ReportSpace::Report external_bug_report;

bool BranchCheckerVisitor::VisitStmt(clang::Stmt *stmt)
{
    switch (stmt->getStmtClass())
    {
    case clang::Stmt::IfStmtClass:
    {
        auto result = _checker.analyzer().results().find(stmt);
        if (result != _checker.analyzer().results().end())
        {
            if (_checker.options().enable_code_layout)
            {
                if (result->second <= _checker.options().unlikely_threshold)
                {
                    const clang::Stmt *then_stmt = llvm::cast<clang::IfStmt>(stmt)->getThen();
                    if (then_stmt->getStmtClass() == clang::Stmt::CompoundStmtClass)
                    {
                        if (llvm::cast<clang::CompoundStmt>(then_stmt)->size() >= _checker.options().code_count_threshold)
                        {
                            generate_report(ReportSpace::DefectName::InstructionLayout,
                                            ReportSpace::DefectType::Warning, then_stmt, "Non-critical code may waste code cache");
                        }
                    }
                }
            }
            if (_checker.options().enable_likely_unlikely)
            {
                if(result->second == 0.0 || result->second == 1.0) // has macro, no detect
                    break;
                if (result->second <= _checker.options().unlikely_threshold)
                {
                    generate_report(ReportSpace::DefectName::AddLikelyOrUnlikelyToBranch,
                                    ReportSpace::DefectType::Warning, stmt, "can use unlikely here");
                }
                else if (result->second >= _checker.options().likely_threshold)
                {
                    generate_report(ReportSpace::DefectName::AddLikelyOrUnlikelyToBranch,
                                    ReportSpace::DefectType::Warning, stmt, "can use likely here");
                }
            }
        }
        break;
    }
    default:
        break;
    }
    return true;
}

void BranchCheckerVisitor::generate_report(ReportSpace::DefectName dn, ReportSpace::DefectType dt, const clang::Stmt *stmt, const std::string &desc)
{
    const clang::SourceManager &source_mgr = _ast_ctx.getSourceManager();
    auto loc = stmt->getBeginLoc();
    unsigned int lineno = source_mgr.getExpansionLineNumber(loc);
    unsigned int colno = source_mgr.getExpansionColumnNumber(loc);
    std::string filename = loc.printToString(source_mgr);
    ReportSpace::Defect *defect = new ReportSpace::Defect(dn,
                                                          dt,
                                                          filename, lineno, colno,
                                                          "");
    defect->addDesc(desc);
    external_bug_report.addToReport(defect);
}

BranchChecker::BranchChecker(ASTResource *resource, ASTManager *manager,
                             CallGraph *call_graph, Config *configure)
    : BasicChecker(resource, manager, call_graph, configure), _analyzer(resource, manager, call_graph, configure)
{
    auto enable = configure->getOptionBlock("CheckerEnable");
    _options.enable_code_layout =
        Config::IsBlockConfigTrue(enable, "InstructionLayoutChecker");
    _options.enable_likely_unlikely = Config::IsBlockConfigTrue(enable, "AddLikelyOrUnlikelyToBranchChecker");
    //
    std::unordered_map<std::string, std::string> code_layout_cfg =
        configure->getOptionBlock("InstructionLayoutChecker");
    auto code_layout_threshold = code_layout_cfg.find("code_count_threshold");
    if (code_layout_threshold != code_layout_cfg.end())
        _options.code_count_threshold = stoi(code_layout_threshold->second);

    std::unordered_map<std::string, std::string> likely_unlikely_cfg =
        configure->getOptionBlock("AddLikelyOrUnlikelyToBranchChecker");
    auto likely_threshold = likely_unlikely_cfg.find("likely_threshold");
    if (likely_threshold != likely_unlikely_cfg.end())
        _options.likely_threshold = stod(likely_threshold->second);
    auto unlikely_threshold = likely_unlikely_cfg.find("unlikely_threshold");
    if (unlikely_threshold != likely_unlikely_cfg.end())
        _options.unlikely_threshold = stod(unlikely_threshold->second);
}

void BranchChecker::check()
{
    /* size_t file_num = 1; */
    size_t file_num = resource->getASTFiles().size();
    size_t processed = 0;
    for (auto ast_file : resource->getASTFiles())
    {
        ++processed;
        process_bar((float)processed / file_num);
        /* llvm::errs() << '[' << file_num++ << '/' << resource->getASTFiles().size() << "] Processing file: " << manager->getASTUnit(ast_file)->getMainFileName().str() << '\n'; */
        _analyzer.analyze(ast_file);
        auto &ctx = manager->getASTUnit(ast_file)->getASTContext();
        BranchCheckerConsumer consumer(*this, ctx);
        consumer.HandleTranslationUnit(ctx);
    }
}